import os
import shutil
from PIL import Image
from tqdm import tqdm

def find_and_remove_corrupted(folder):
    corrupted = []
    quarantine = os.path.join(folder, "_corrupted")
    os.makedirs(quarantine, exist_ok=True)
    
    for root, _, files in os.walk(folder):
        if "_corrupted" in root:
            continue
            
        for file in tqdm(files, desc=f"扫描 {root}"):
            if file.lower().endswith(('.png', '.jpg', '.jpeg')):
                path = os.path.join(root, file)
                try:
                    with Image.open(path) as img:
                        img.verify()
                    with Image.open(path) as img:
                        img.load()  # 强制完全加载
                except Exception as e:
                    corrupted.append(path)
                    print(f"\n发现损坏文件: {path}")
                    try:
                        shutil.move(path, os.path.join(quarantine, file))
                    except Exception as move_error:
                        print(f"移动失败: {move_error}")
                        os.remove(path)  # 强制删除
    return corrupted

if __name__ == "__main__":
    data_dir = "E:/sdxxproject/data"
    print("开始深度扫描数据集...")
    bad_files = find_and_remove_corrupted(data_dir)
    
    print(f"\n处理完成！共发现并处理 {len(bad_files)} 个损坏文件")
    print(f"已移动到: {os.path.join(data_dir, '_corrupted')}")